package RBT;

import BST.BsTreeNode;

import static BST.BSTUtils.searchRecursive;

/**
 * 红黑树
 *
 * @author Lwh
 * @time 2022/11/11 02:15
 */
public class RbTree<T extends Comparable<T>> {
    /*根节点*/
    private RbTreeNode<T> mRoot;
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RbTree() {
        mRoot = null;
    }

    /*获取某个节点的父节点*/
    public RbTreeNode<T> parentOf(RbTreeNode<T> node) {
        return node != null ? node.parent : null;
    }

    /*获取某个节点的颜色*/
    public boolean colorOf(RbTreeNode<T> node) {
        return node != null ? node.color : BLACK;
    }

    /*获取某个树的根节点*/
    public RbTreeNode<T> getMRoot() {
        return this.mRoot;
    }

    /*判断某个节点是否为红色*/
    public boolean isRed(RbTreeNode<T> node) {
        return (node != null) && (node.color == RED);
    }

    /*判断某个节点是否为黑色*/
    public boolean isBlack(RbTreeNode<T> node) {
        return !isRed(node);
    }

    /*设置某个节点为黑色*/
    public void setBlack(RbTreeNode<T> node) {
        if (node != null) {
            node.color = BLACK;
        }
    }

    /*设置某个节点为红色*/
    public void setRed(RbTreeNode<T> node) {
        if (node != null) {
            node.color = RED;
        }
    }

    /*设置某个节点的父节点*/
    public void setParent(RbTreeNode<T> node, RbTreeNode<T> parent) {
        if (node != null) {
            node.parent = parent;
        }
    }

    /*设置某个节点的颜色*/
    public void setColor(RbTreeNode<T> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    /**
     * <pre>
     * 左旋<br>
     *     P                V
     *    / \              / \
     *   F   V     ->     P   X
     *      / \          / \
     *     R   X        F   R
     * </pre>
     */
    public void leftRotate(RbTreeNode<T> pNode) {
        //1.P的右子节点V的左子节点R变为旋转节点P的右子节点
        RbTreeNode<T> vNode = pNode.right;
        RbTreeNode<T> rNode = vNode.left;
        pNode.right = rNode;

        if (rNode != null) {
            //2.修正R的父节点为P
            rNode.parent = pNode;
        }

        //3.修正V的父节点为P原来的父节点
        vNode.parent = pNode.parent;

        if (pNode.parent == null) {
            //如果P原来没有父节点，说明为根节点
            this.mRoot = vNode;
        } else {
            //4.如果P原来有父节点，则V取代P作为这个父节点的左孩子或右孩子
            if (pNode.parent.left == pNode) {
                pNode.parent.left = vNode;
            } else {
                pNode.parent.right = vNode;
            }
        }
        vNode.left = pNode;
        pNode.parent = vNode;
    }

    /**
     * <pre>
     * 右旋<br>
     *       P              F
     *      / \            / \
     *     F   V    ->    D   P
     *    / \                / \
     *   D   K              K   V
     * </pre>
     */
    public void rightRotate(RbTreeNode<T> pNode) {
        RbTreeNode<T> fNode = pNode.left;
        RbTreeNode<T> kNode = fNode.right;
        //移动K
        pNode.left = kNode;
        //修正K的父节点
        if (kNode != null) {
            kNode.parent = pNode;
        }
        //修正F的父节点
        if (pNode.parent == null) {
            this.mRoot = fNode;
        } else {
            if (pNode == pNode.parent.left) {
                pNode.parent.left = fNode;
            } else {
                pNode.parent.right = fNode;
            }
        }
        //调整P和F的父子关系
        fNode.right = pNode;
        pNode.parent = fNode;
    }

    /**
     * 递归实现红黑树插入
     *
     * @param root         根节点引用
     * @param nodeInserted 待插入节点
     */
    public boolean insertRecursively(RbTreeNode<T> root,
                                     RbTreeNode<T> nodeInserted) {
        if (root == null) {
            mRoot = nodeInserted;
            return true;
        }

        if (root.nodeKey.compareTo(nodeInserted.nodeKey) > 0) {
            if (root.left == null) {
                //1.将红黑树当作一颗二叉查找树，将节点添加到二叉查找树中
                root.left = nodeInserted;
                nodeInserted.parent = root;
                //2.设置插入节点的颜色为红色
                nodeInserted.color = RED;
                //3.修正
                insertFixUp(nodeInserted);
                return true;
            } else {
                //向左子树中递归操作
                return insertRecursively(root.left, nodeInserted);
            }
        } else if (root.nodeKey.compareTo(nodeInserted.nodeKey) < 0) {
            if (root.right == null) {
                //1.将红黑树当作一颗二叉查找树，将节点添加到二叉查找树中
                root.right = nodeInserted;
                nodeInserted.parent = root;
                //2.设置插入节点的颜色为红色
                nodeInserted.color = RED;
                //3.修正
                insertFixUp(nodeInserted);
                return true;
            } else {
                //向右子树中递归操作
                return insertRecursively(root.right, nodeInserted);
            }
        } else {
            return false;
        }
    }

    /**
     * 插入操作修复函数
     */
    private void insertFixUp(RbTreeNode<T> currentNode) {
        RbTreeNode<T> parent;
        RbTreeNode<T> gParent;
        //insert-fixup-case-1:被插入节点是根节点，将根节点设为黑色
        if (currentNode == mRoot) {
            setBlack(this.mRoot);
            return;
        }
        //insert-fixup-case-2:被插入节点的父节点是黑色，什么都不需要做
        if (parentOf(currentNode) != null && isBlack(parentOf(currentNode))) {
            //do nothing
            return;
        }
        //insert-fixup-case-3:当前节点的父节点存在，并且父节点的颜色是红色
        while (((parent = parentOf(currentNode)) != null) && isRed(parent)) {
            gParent = parentOf(parent);
            //若当前节点的父节点是祖父节点的左孩子
            if (parent == gParent.left) {
                //insert-fixup-case-3.1:当前节点的叔叔节点是红色
                RbTreeNode<T> uncle = gParent.right;
                if ((uncle != null) && isRed(uncle)) {
                    setBlack(uncle);
                    setBlack(parent);
                    setRed(gParent);
                    currentNode = gParent;
                    continue;
                }
                //insert-fixup-case-3.2:当前节点的叔叔节点是黑色，且当前节点是右孩子
                if (parent.right == currentNode) {
                    RbTreeNode<T> tmp;
                    leftRotate(parent);
                    tmp = parent;
                    parent = currentNode;
                    currentNode = tmp;
                }
                //insert-fixup-case-3.3:当前节点的叔叔节点是黑色，且当前节点是左孩子
                setBlack(parent);
                setRed(gParent);
                rightRotate(gParent);
            }
            //若当前节点的父节点是祖父节点的右孩子
            else {
                RbTreeNode<T> uncle = gParent.left;
                if ((uncle != null) && isRed(uncle)) {
                    setBlack(uncle);
                    setBlack(parent);
                    setRed(gParent);
                    currentNode = gParent;
                    continue;
                }
                if (parent.left == currentNode) {
                    RbTreeNode<T> tmp;
                    rightRotate(parent);
                    tmp = parent;
                    parent = currentNode;
                    currentNode = tmp;
                }
                setBlack(parent);
                setRed(gParent);
                leftRotate(gParent);
            }
        }
    }

    /**
     * 删除键值为key的节点
     *
     * @param key 待删除的节点值
     */
    public void remove(T key) {
        BsTreeNode<T> node;
        if ((node = searchRecursive(mRoot, key)) != null) {
            remove(node.getValue());
        }
    }

    /**
     * 删除节点
     *
     * @param nodeRemoved 待删除节点
     */
    public void remove(RbTreeNode<T> nodeRemoved) {
        //被删除节点的左右孩子都不为空
        if ((nodeRemoved.left != null) && (nodeRemoved.right != null)) {
            removeNodeWithDoubleChild(nodeRemoved);
        }
        //被删除节点的左孩子不为空
        else if (nodeRemoved.left != null) {
            removeNodeWithOnlyLeftChild(nodeRemoved);
        }
        //被删除节点的右孩子不为空
        else if (nodeRemoved.right != null) {
            removeNodeWithOnlyRightChild(nodeRemoved);
        }
        //被删除节点的左右孩子都为空
        else {
            removeNodeWithNoChild(nodeRemoved);
        }
    }

    private void removeNodeWithNoChild(RbTreeNode<T> nodeRemoved) {
        RbTreeNode<T> childOfNodeRemoved;
        RbTreeNode<T> parentOfNodeRemoved;
        boolean colorOfNodeRemoved;

        childOfNodeRemoved = null;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;

        //被删除节点不是根节点（根节点不存在父节点）
        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
            //被删除节点是根节点，更新根节点
            this.mRoot = childOfNodeRemoved;
        }

        //删除修正
        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    private void removeNodeWithOnlyLeftChild(RbTreeNode<T> nodeRemoved) {
        RbTreeNode<T> childOfNodeRemoved;
        RbTreeNode<T> parentOfNodeRemoved;
        boolean colorOfNodeRemoved;

        childOfNodeRemoved = nodeRemoved.left;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;

        childOfNodeRemoved.parent = parentOfNodeRemoved;

        //被删除节点不是根节点（根节点不存在父节点）
        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
            //被删除节点是根节点，更新根节点
            this.mRoot = childOfNodeRemoved;
        }

        //删除修正
        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    private void removeNodeWithOnlyRightChild(RbTreeNode<T> nodeRemoved) {
        RbTreeNode<T> childOfNodeRemoved;
        RbTreeNode<T> parentOfNodeRemoved;
        boolean colorOfNodeRemoved;

        childOfNodeRemoved = nodeRemoved.right;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;

        childOfNodeRemoved.parent = parentOfNodeRemoved;

        //被删除节点不是根节点（根节点不存在父节点）
        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
            //被删除节点是根节点，更新根节点
            this.mRoot = childOfNodeRemoved;
        }

        //删除修正
        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    private void removeNodeWithDoubleChild(RbTreeNode<T> nodeRemoved) {
        //找到当前节点的中序后继节点，成为取代节点
        RbTreeNode<T> replace = nodeRemoved.right;
        while (replace.left != null) {
            replace = replace.left;
        }

        /*
         * 拷贝取代节点的内容到被删除节点
         * （这里值拷贝nodeKey，实际上一个节点除了nodeKey域还有nodeData域,只不过我们简化了节点定义）
         */
        nodeRemoved.nodeKey = replace.nodeKey;
        //然后删除替代节点
        remove(replace);
    }

    /**
     * 删除操作修复函数
     */
    private void removeFixUp(RbTreeNode<T> childOf,
                             RbTreeNode<T> parentOf) {
        //other是childOf的兄弟节点
        RbTreeNode<T> other;

        //remove-fixup-case-1:childOf是“红+黑”节点
        //直接把childOf设为黑色
        if (childOf != null && isRed(childOf)) {
            setBlack(childOf);
        }
        //remove-fixup-case-2:childOf是“黑+黑”节点
        //且childOf是根，什么都不做
        if (childOf != null && isBlack(childOf) && (childOf == this.mRoot)) {
            //do nothing
            return;
        }

        //remove-fixup-case-3:childOf是“黑+黑”节点
        //且childOf不是根节点
        while ((childOf == null || isBlack(childOf)) && (childOf != this.mRoot)) {
            if (parentOf.left == childOf) {
                other = parentOf.right;
                if (isRed(other)) {
                    //remove-fixup-case-3.1:childOf的兄弟 other是红色
                    setBlack(other);
                    setRed(parentOf);
                    leftRotate(parentOf);
                    other = parentOf.right;
                }
                if ((other.left == null || isBlack(other.left)) && (other.right == null || isBlack(other.right))) {
                    //remove-fixup-case-3.2:childOf的兄弟 other是黑色，且other的两个孩子也都是黑色
                    setRed(other);
                    childOf = parentOf;
                    parentOf = parentOf(childOf);
                } else {
                    if (other.right == null || isBlack(other.right)) {
                        //remove-fixup-case-3.3:childOf的兄弟 other是黑色，且other的左孩子是红色，右孩子是黑色
                        setBlack(other.left);
                        setRed(other);
                        rightRotate(other);
                        other = parentOf.right;
                    }
                    //remove-fixup-case-3.4:childOf的兄弟 other是黑色，且other的右孩子是红色，左孩子是任意颜色
                    setColor(other, colorOf(parentOf));
                    setBlack(parentOf);
                    setBlack(other.right);
                    leftRotate(parentOf);
                    childOf = this.mRoot;
                    break;
                }
            } else {
                other = parentOf.left;
                if (isRed(other)) {
                    //3.1
                    setBlack(other);
                    setRed(parentOf);
                    rightRotate(parentOf);
                    other = parentOf.left;
                }
                if ((other.left == null || isBlack(other.left)) && (other.right == null || isBlack(other.right))) {
                    //3.2
                    setRed(other);
                    childOf = parentOf;
                    parentOf = parentOf(childOf);
                } else {
                    if (other.left == null || isBlack(other.left)) {
                        //3.3
                        setBlack(other.right);
                        setRed(other);
                        leftRotate(other);
                        other = parentOf.left;
                    }
                    //3.4
                    setColor(other, colorOf(parentOf));
                    setBlack(parentOf);
                    setBlack(other.left);
                    rightRotate(parentOf);
                    childOf = this.mRoot;
                    break;
                }
            }
        }
    }
}
